{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCP Proxy\n",
    "\n",
    "The [Model Context Protocol (MCP)](https://modelcontextprotocol.io/introduction) is an open protocol that defines a standardized way to provide context to large language models (LLMs). By using MCP, applications can seamlessly interact with LLMs in a structured and efficient manner, making it easier to exchange information, use external tools, and enhance the capabilities of language models in a consistent and reusable way.\n",
    "\n",
    "This guide will walk you through how to create **MCP clients** that connect to MCP servers, enabling you to utilize external tools and processes efficiently.\n",
    "\n",
    "## Introduction to MCP\n",
    "\n",
    "MCP aims to standardize how applications provide contextual information to LLMs and how LLMs can access external tools. MCP servers expose specific tools that are accessible via different communication protocols, such as [`stdio` (Standard Input/Output)](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio), [`SSE` (Server-Sent Events)](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse), and [`streamable-http`(Streamable HTTP)](https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#streamable-http).\n",
    "\n",
    "MCP clients are responsible for interacting with these servers, sending requests, and processing responses. By integrating MCP into your application, you can build systems where LLMs can utilize external computational tools in real-time.\n",
    "\n",
    "In this guide, we will:\n",
    "\n",
    "- Set up an MCP server that exposes basic mathematical tools (`add` and `multiply`).\n",
    "- Connect the AG2 framework to this MCP server.\n",
    "- Communicate using two different transport protocols: [stdio](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio) and [SSE](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse).\n",
    "\n",
    "## Installation\n",
    "\n",
    "To integrate MCP tools into the AG2 framework, install the required dependencies:\n",
    "\n",
    "```bash\n",
    "pip install -U \"ag2[openai,mcp,mcp-proxy-gen]>=0.9.4\"\n",
    "```\n",
    "\n",
    "> **Note:** If you are using `pyautogen`, uninstall it before installing `ag2` to ensure compatibility:\n",
    "> ```bash\n",
    "> pip uninstall -y pyautogen\n",
    "> ```\n",
    "> ```bash\n",
    "> pip install -U \"ag2[openai,mcp,mcp-proxy-gen]>=0.9.4\"\n",
    "> ``` \n",
    "\n",
    "## Imports\n",
    "\n",
    "Before diving into the code, let’s go over the main imports used in this guide:\n",
    "\n",
    "- `ClientSession`: Manages the client session for connecting to the MCP server, allowing you to send requests and handle responses.\n",
    "- `StdioServerParameters`: Provides parameters to set up communication with an MCP server over `stdio`.\n",
    "- `stdio_client`: A utility that helps you connect to an MCP server using the standard input/output protocol.\n",
    "- `sse_client`: A utility to connect to an MCP server using the `SSE` (Server-Sent Events) protocol.\n",
    "- `create_toolkit`: A helper function that creates a toolkit, wrapping the tools exposed by the MCP server for easier access.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You also need an email validator for this example:\n",
    "    \n",
    "`pip install \"pydantic[email]\"`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import types\n",
    "\n",
    "# Create a dummy module to fake out the import\n",
    "dummy_module = types.ModuleType(\"patches\")\n",
    "dummy_module.patch_parse = lambda: None\n",
    "\n",
    "# Inject it into sys.modules *before* importing the library\n",
    "sys.modules[\"fastapi_code_generator.patches\"] = dummy_module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import timedelta\n",
    "from pathlib import Path\n",
    "\n",
    "from mcp import ClientSession, StdioServerParameters\n",
    "from mcp.client.stdio import stdio_client\n",
    "\n",
    "from autogen import LLMConfig\n",
    "from autogen.agentchat import AssistantAgent\n",
    "from autogen.mcp import create_toolkit\n",
    "from autogen.mcp.mcp_proxy import MCPProxy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only needed for Jupyter notebooks\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Define our LLM configuration for OpenAI's GPT-4o mini,\n",
    "#    uses the OPENAI_API_KEY environment variable\n",
    "# llm_config = LLMConfig(config_list={\"api_type\": \"openai\", \"model\": \"gpt-5-nano\"})\n",
    "llm_config = LLMConfig.from_json(path=\"../OAI_CONFIG_LIST\").where(model=\"gpt-5-nano\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Proxy MCP Server Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "tmp_path = Path(\"tmp\") / \"mcp_trello\"\n",
    "shutil.rmtree(tmp_path, ignore_errors=True)\n",
    "tmp_path.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Create an MCP server code via the MCPProxy (requires dependencies from mcp-proxy-gen option)\n",
    "MCPProxy.create(\n",
    "    openapi_url=\"https://raw.githubusercontent.com/Asana/openapi/master/defs/asana_oas.yaml\",\n",
    "    client_source_path=tmp_path,\n",
    "    rename_functions=True,\n",
    "    group_functions=True,\n",
    "    configuration_type=\"yaml\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "assert tmp_path.exists(), \"Failed to create tmp directory\"\n",
    "# assert list(tmp_path.glob(\"*.py\")) >= 2, \"Failed to generate OpenAPI client files\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a Toolkit from MCP Tools\n",
    "\n",
    "To allow AG2 to interact with the MCP server, we need to establish a connection, create a toolkit from the tools exposed by the server, and then register this toolkit with an AG2 agent. The toolkit will provide the necessary functionality for AG2 to call the MCP tools, such as performing mathematical operations.\n",
    "\n",
    "**Steps to Create a Toolkit from MCP Tools**\n",
    "\n",
    "- **Wrap MCP tools into a toolkit**: The tools exposed by the MCP server (like `add` and `multiply`) are wrapped into a toolkit for easy use by AG2.\n",
    "- **Register the toolkit with an AG2 agent**: This step allows AG2 to use the toolkit’s tools when processing requests.\n",
    "- **Start a conversation**: AG2 can now send a message to the MCP server requesting a task to be performed (e.g., adding two numbers) and receive the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def create_toolkit_and_run(session: ClientSession) -> None:\n",
    "    # Create a toolkit with available MCP tools\n",
    "    toolkit = await create_toolkit(session=session)\n",
    "    agent = AssistantAgent(\n",
    "        name=\"assistant\", llm_config=LLMConfig(config_list={\"model\": \"gpt-5-nano\", \"api_type\": \"openai\"})\n",
    "    )\n",
    "    # Register MCP tools with the agent\n",
    "    toolkit.register_for_llm(agent)\n",
    "\n",
    "    # Make a request using the MCP tool\n",
    "    result = await agent.a_run(\n",
    "        message=\"List all the asana functions that you can use\",\n",
    "        tools=toolkit.tools,\n",
    "        max_turns=2,\n",
    "        user_input=False,\n",
    "    )\n",
    "\n",
    "    await result.process()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting Up the Client Session and Communicating with the MCP Server\n",
    "\n",
    "Once the toolkit is created, the next step is to connect AG2 to the MCP server and manage the communication between them. This involves setting up connection parameters and starting a client session, which will facilitate the entire interaction. After establishing the connection, AG2 can send requests and receive responses from the server.\n",
    "\n",
    "### Steps to Set Up the Client Session and Communicate with the MCP Server\n",
    "\n",
    "- **Set up connection parameters**: Define how AG2 will communicate with the MCP server. This includes specifying the transport protocol (`stdio` or `SSE`) and other connection details.\n",
    "- **Start a client session**: This session facilitates communication between AG2 and the MCP server. The session is responsible for sending requests, receiving responses, and maintaining the connection state.\n",
    "\n",
    "#### Using `stdio_client` for Communication\n",
    "\n",
    "Establish the connection in the client code by using the `stdio_client` requires:\n",
    "\n",
    "- **Server Parameters**: The connection is defined using the `StdioServerParameters`, where the command (`python`) runs the server (`math_server.py`) with `stdio` mode.\n",
    "- **Connecting to the Server**: The `stdio_client` establishes the connection to the server, and the `ClientSession` is used to facilitate communication.\n",
    "- **Running the Task**: Once the connection is established, `create_toolkit_and_run(session)` is called to wrap the MCP tools and perform the task asynchronously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create server parameters for stdio connection\n",
    "\n",
    "# security = APIKeyHeader.Parameters(value=\"some_api_key\")\n",
    "\n",
    "server_params = StdioServerParameters(\n",
    "    # add security parameters here as env var\n",
    "    command=\"python\",  # The command to run the server\n",
    "    args=[\n",
    "        str(tmp_path / \"main.py\"),  # Path to the server script\n",
    "        \"stdio\",\n",
    "    ],  # Path to server script and transport mode\n",
    "    env={\n",
    "        \"CONFIG_PATH\": str(\n",
    "            tmp_path.absolute() / \"mcp_config_member_management.json\"\n",
    "        ),  # Path to the config file, check the generated files as the names change with each generation\n",
    "    },\n",
    "    # env={\n",
    "    #     \"SECURITY\": str(security.dump()),\n",
    "    # },\n",
    ")\n",
    "\n",
    "async with (\n",
    "    stdio_client(server_params) as (read, write),\n",
    "    ClientSession(read, write, read_timeout_seconds=timedelta(seconds=30)) as session,\n",
    "):\n",
    "    # Initialize the connection\n",
    "    await session.initialize()\n",
    "    await create_toolkit_and_run(session)"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "MCP Clients",
   "tags": [
    "MCP",
    "Model Context Protocol",
    "tools"
   ]
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
